import time

from torch.utils.data import DataLoader,Dataset
import torch
import numpy as np
import albumentations as A
from albumentations.pytorch import ToTensorV2
from torchvision import transforms
train_transformer = A.Compose([
    A.Resize(256,256),
    A.OneOf([
        A.Rotate(90,90,p=0.5),
    ]),
    #A.ToGray(),
    #A.Normalize(),
    ToTensorV2(),
])
# transform = A.Compose([
#     A.ToFloat(max_value=65535.0),
#     A.RandomRotate90(),
#     A.Flip(),
#     A.OneOf([
#         A.MotionBlur(p=0.2),
#         A.MedianBlur(blur_limit=3, p=0.1),
#         A.Blur(blur_limit=3, p=0.1),
#     ], p=0.2),
#     A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.2, rotate_limit=45, p=0.2),
#     A.OneOf([
#         A.OpticalDistortion(p=0.3),
#         A.GridDistortion(p=0.1),
#     ], p=0.2),
#     A.HueSaturationValue(hue_shift_limit=20, sat_shift_limit=0.1, val_shift_limit=0.1, p=0.3),
#
#     A.FromFloat(max_value=65535.0),
# ])
class Mydataset(Dataset):
    def __init__(self,train_input_common,train_gt_common,transformer=None,mode='train'):
        self.train_input_common = train_input_common
        self.train_gt_common = train_gt_common
        self.transformer = transformer
        self.mode = mode
        self.sample_ref = np.frombuffer(train_input_common, dtype='uint16').reshape((-1, 256, 256))
        self.sample_gt = np.frombuffer(train_gt_common, dtype='uint16').reshape((-1, 256, 256))
    def __len__(self):
        return len(self.sample_ref)
    def __getitem__(self, item):
        if self.mode == 'train':
            input = self.sample_ref[item].astype(float)* np.float32(1 / 65536)
            input = np.expand_dims(input,axis=2)

            output = self.sample_gt[item].astype(float)* np.float32(1 / 65536)
            #output = np.expand_dims(output,axis=2)

            augmented = self.transformer(image = input,mask = output)
            return augmented['image'].to(torch.float32),augmented['mask'].unsqueeze(0).to(torch.float32)
            #return input


if __name__ == '__main__':
    train_input_path = '/Users/lihaobo/Downloads/burst_raw/competition_train_input.0.2.bin'
    train_gt_path = '/Users/lihaobo/Downloads/burst_raw/competition_train_gt.0.2.bin'
    input_common = open(train_input_path, 'rb').read()
    gt_common = open(train_gt_path, 'rb').read()
    dataset = Mydataset(train_input_common= input_common,train_gt_common =gt_common,transformer=train_transformer)
    #print(dataset[0][0].shape)
    #start = time.time()
    dataloader = DataLoader(dataset,batch_size=128)
    for i,(input,output) in enumerate(dataloader):
        if i==0:
            start = time.time()
            print(input.shape,output.shape)

        print(i)
    print("一轮时间{0}".format(time.time()-start))